import json
from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载微调后的模型和 tokenizer
def load_finetuned_model(model_dir):
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_pretrained(model_dir)
    return tokenizer, model

# 生成中间代码
def generate_middle_code(tokenizer, model, prefix, fim_suffix, max_length=512):
    # 组合输入
    input_text = f"{prefix} <MIDDLE> {fim_suffix}"
    inputs = tokenizer(input_text, return_tensors="pt")

    # 生成代码
    outputs = model.generate(
        inputs["input_ids"],
        max_length=max_length,
        num_return_sequences=1,  # 生成一个结果
        pad_token_id=tokenizer.eos_token_id,  # 设置结束符
        do_sample=True,  # 启用采样
        top_k=50,  # 限制采样范围
        top_p=0.95,  # Nucleus 采样
        temperature=0.7,  # 控制随机性
    )

    # 解码生成的代码
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # 提取中间部分
    middle_start = generated_text.find(prefix) + len(prefix)
    middle_end = generated_text.find(fim_suffix)
    middle_code = generated_text[middle_start:middle_end].strip()

    return middle_code

def main():
    # 微调后的模型路径
    model_dir = "./codellama-7b-finetuned"

    # 加载模型和 tokenizer
    tokenizer, model = load_finetuned_model(model_dir)

    # 输入 JSON 文件路径
    input_json_path = "Q_A.json"
    # 输出 JSON 文件路径
    output_json_path = "output.json"

    # 读取输入 JSON 文件
    with open(input_json_path, "r", encoding="utf-8") as f:
        data = [json.loads(line) for line in f]

    # 生成中间代码并保存结果
    results = []
    for item in data:
        prefix = item["prefix"]
        fim_suffix = item["fim_suffix"]

        # 生成中间代码
        middle_code = generate_middle_code(tokenizer, model, prefix, fim_suffix)

        # 保存结果
        result = {
            "prefix": prefix,
            "middle": middle_code,
            "fim_suffix": fim_suffix
        }
        results.append(result)

    # 将结果保存到输出 JSON 文件
    with open(output_json_path, "w", encoding="utf-8") as f:
        for result in results:
            f.write(json.dumps(result, ensure_ascii=False) + "\n")

if __name__ == "__main__":
    main()